import numpy as np
import torch

from bert_yt.b_config import b_conf


class OnnxModelRunner:
    """
    ONNX模型推理器，用于实际部署场景
    """

    def __init__(self, model_path=None, pretrained_model_path=None):
        """
        初始化ONNX模型推理器
        :param model_path: ONNX模型路径，默认使用配置中的路径
        :param pretrained_model_path: 预训练模型路径，用于加载分词器
        """
        import onnxruntime as ort
        from transformers import BertTokenizer
        from utils import id2label

        if model_path is None:
            model_path = b_conf.onnx_model_path

        if pretrained_model_path is None:
            pretrained_model_path = b_conf.pretrained_model_path

        # 创建一次session对象，后续重复使用（重要：避免重复创建）
        self.session = ort.InferenceSession(model_path)
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_path)
        self.id2label = id2label

        # 获取模型输入输出信息
        self.input_names = [input.name for input in self.session.get_inputs()]
        self.output_names = [output.name for output in self.session.get_outputs()]

    def predict(self, texts):
        """
        对文本进行推理预测
        :param texts: 单个文本字符串或文本列表
        :return: 预测标签列表
        """
        # 确保输入是列表格式
        if isinstance(texts, str):
            texts = [texts]

        # 对文本进行预处理
        tokens = self.tokenizer.batch_encode_plus(
            texts,
            return_tensors='np',  # 使用numpy数组以适配ONNX Runtime
            max_length=b_conf.max_length,
            padding=True,
            truncation=True,
        )

        input_ids = tokens['input_ids'].astype(np.int64)
        attention_mask = tokens['attention_mask'].astype(np.int64)

        # 运行推理（重复使用已创建的session）
        outputs = self.session.run(
            None,  # 获取所有输出
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask
            }
        )

        # 处理输出结果
        predictions = outputs[0]  # 获取第一个输出（logits）
        predicted_classes = predictions.argmax(axis=1)  # 获取预测类别索引
        predicted_labels = [self.id2label[class_id] for class_id in predicted_classes]

        return predicted_labels


# 全局模型实例，用于实际部署（避免重复加载模型）
_model_runner = None


def get_model_runner():
    """
    获取全局模型推理器实例（单例模式）
    在实际部署中，推荐使用这种方式避免重复加载模型
    """
    global _model_runner
    if _model_runner is None:
        _model_runner = OnnxModelRunner()
    return _model_runner


def to_onnx():
    from bert_yt.my_bert import get_model
    model = get_model()
    model.eval()  # 设置模型为评估模式
    # 示例输入，用于推断模型结构、验证模型逻辑
    # 形状为 [batch_size, sequence_length]
    input_tensor = torch.torch.randint(0, 1000, (1, 512)).long()
    attention_mask = torch.randn((1, 512)).long()
    torch.onnx.export(
        model,  # model to export
        (input_tensor, attention_mask),  # inputs of the model,
        b_conf.onnx_model_path,  # filename of the ONNX model
        input_names=["input_ids", "attention_mask"],  # Rename inputs for the ONNX model
        output_names=["output"],  # 显式命名输出节点为 "output"，便于后续推理时直接调用。
        opset_version=14,  # 显式指定 opset_version=13，确保兼容性和稳定性。
        do_constant_folding=True,  # 优化常量折叠，减少模型计算图复杂度
        dynamic_axes={
            "input_ids": {0: "batch_size"},  # 可变的batch_size维度
            "attention_mask": {0: "batch_size"},  # 可变的batch_size维度
            "output": {0: "batch_size"}  # 可变的batch_size维度
        }
    )
    print("导出 ONNX 模型成功！")


def load_onnx():
    import onnxruntime as ort

    # 加载ONNX模型
    session = ort.InferenceSession(b_conf.onnx_model_path)

    # 获取模型的输入和输出名称
    inputs = session.get_inputs()
    outputs = session.get_outputs()

    print("Model inputs:")
    for input in inputs:
        print(f"  Name: {input.name}, Shape: {input.shape}, Type: {input.type}")

    print("Model outputs:")
    for output in outputs:
        print(f"  Name: {output.name}, Shape: {output.shape}, Type: {output.type}")


def load_and_inference(sample_texts=None):
    """
    加载ONNX模型，并对实际数据进行推理
    :param sample_texts: 样本文本列表，如果为None则使用配置中的测试文本
    """
    from config import conf

    # 如果没有提供样本文本，则使用配置中的测试文本
    if sample_texts is None:
        sample_texts = conf.test_texts[:3]  # 使用前3个测试文本作为示例
    runner = get_model_runner()
    predicted_classes = runner.predict(sample_texts)

    # 打印结果
    print("ONNX模型推理结果:")
    print("=" * 50)
    for i, (text, pred_class) in enumerate(zip(sample_texts, predicted_classes)):
        print(f"文本 {i + 1}: {text[:100]}...")
        print(f"预测类别: {pred_class}")
        print("-" * 30)

    return predicted_classes


if __name__ == '__main__':
    # to_onnx()
    # load_onnx()
    load_and_inference()
    pass
